#include "kernel_operator.h"
#include "op_common.h"
#include <cstdint>

using namespace AscendC;

template <typename _DT_XYZ1, typename _DT_XYZ2, typename _DT_DIST,
          typename _DT_INDICES>
class Kernel {
public:
  TPipe pipe;
  DefGlobalTensor(XYZ1);
  DefGlobalTensor(XYZ2);
  DefGlobalTensor(DIST);
  DefGlobalTensor(INDICES);
  int64_t b = 0;
  int64_t n = 0;
  int64_t m = 0;
  int64_t tileLength = 0;
  int64_t tileVectorSize = 0;
  DefInQue(XYZ1);
  DefInQue(XYZ2);
  DefBufVECCALC(TEMP);
  DefBufVECCALC(WORK_LOCAL);

public:
  __aicore__ inline Kernel() {}

  template <typename T>
  __aicore__ inline void Init(GM_ADDR xyz1, GM_ADDR xyz2, GM_ADDR dist,
                              GM_ADDR indices, T &tiling_data) {
    b = tiling_data.b;
    n = tiling_data.n;
    m = tiling_data.m;
    tileLength = tiling_data.tileLength;
    tileVectorSize = tileLength / elementsPerBlock<_DT_XYZ2>();
    auto xyz1_size = b * n * 3;
    auto xyz2_size = b * m * 3;
    // 切分global
    GSetBuffer(XYZ1, xyz1, 0, xyz1_size);
    GSetBuffer(XYZ2, xyz2, 0, xyz2_size);
    GSetBuffer(DIST, dist, 0, xyz1_size);
    GSetBuffer(INDICES, indices, 0, xyz1_size);

    InitQueueSimple(XYZ1, 256);
    InitQueueSimple(XYZ2, tileLength);
    InitTBufBuffer(TEMP, tileLength * sizeof(_DT_XYZ2));
    InitTBufBuffer(WORK_LOCAL, tileLength * sizeof(_DT_XYZ2));
  }

  __aicore__ inline void Process() {
    auto offset_xyz1_1 = n * 3;
    auto offset_xyz2_1 = m * 3;
    for (auto bi = 0; bi < b; ++bi) {
      for (auto ni = 0; ni < n; ++ni) {
        auto index_1 = bi * offset_xyz1_1 + ni * 3;
        TypeOf(XYZ1) best1 = 1e40, best2 = 1e40, best3 = 1e40;
        int32_t besti1 = 0, besti2 = 0, besti3 = 0;
        for (auto mi = 0; mi < m; ++mi) {
          auto index_2 = bi * offset_xyz2_1 + mi * 3;
          auto d1 = GTensorName(XYZ2).GetValue(index_2) -
                    GTensorName(XYZ1).GetValue(index_1);
          auto d2 = GTensorName(XYZ2).GetValue(index_2 + 1) -
                    GTensorName(XYZ1).GetValue(index_1 + 1);
          auto d3 = GTensorName(XYZ2).GetValue(index_2 + 2) -
                    GTensorName(XYZ1).GetValue(index_1 + 2);
          d1 = d1 * d1 + d2 * d2 + d3 * d3;
          if (d1 < best1) {
            best3 = best2;
            besti3 = besti2;
            best2 = best1;
            besti2 = besti1;
            best1 = d1;
            besti1 = mi;
          } else if (d1 < best2) {
            best3 = best2;
            besti3 = besti2;
            best2 = d1;
            besti2 = mi;
          } else if (d1 < best3) {
            best3 = d1;
            besti3 = mi;
          }
        }
        GTensorName(DIST).SetValue(index_1, best1);
        GTensorName(DIST).SetValue(index_1 + 1, best2);
        GTensorName(DIST).SetValue(index_1 + 2, best3);
        GTensorName(INDICES).SetValue(index_1, besti1);
        GTensorName(INDICES).SetValue(index_1 + 1, besti2);
        GTensorName(INDICES).SetValue(index_1 + 2, besti3);
      }
    }
  }
  __aicore__ inline int64_t copyIn(int64_t bi, int64_t ni, int64_t mi) {
    auto xyz2Offset = bi * m * 3 + mi * 3;
    auto vectorCount = GET_MIN(tileVectorSize, m - mi);
    QueAllocSimple(XYZ2);
    for (auto i = 0; i < vectorCount; ++i) {
      DataCopy(LTensor(XYZ2)[i * elementsPerBlock<_DT_XYZ2>()],
               GTensor(XYZ2)[xyz2Offset + i * 3], elementsPerBlock<_DT_XYZ2>());
    }
    EnQue(XYZ2);
    return vectorCount;
  }
  __aicore__ inline void tryBest(_DT_XYZ1 v, _DT_INDICES vi, _DT_XYZ1 &dist1,
                                 _DT_XYZ1 &dist2, _DT_XYZ1 &dist3,
                                 _DT_INDICES &idx1, _DT_INDICES &idx2,
                                 _DT_INDICES &idx3) {
    if (v < dist1) {
      dist3 = dist2;
      idx3 = idx2;
      dist2 = dist1;
      idx2 = idx1;
      dist1 = v;
      idx1 = vi;
    } else if (v < dist2) {
      dist3 = dist2;
      idx3 = idx2;
      dist2 = v;
      idx2 = vi;
    } else if (v < dist3) {
      dist3 = v;
      idx3 = vi;
    }
  }

  __aicore__ inline void compute(LocalTensor<_DT_XYZ1> xyz1, int64_t bi,
                                 int64_t ni, int64_t mi, int64_t vectorSize,
                                 _DT_XYZ1 &dist1, _DT_XYZ1 &dist2,
                                 _DT_XYZ1 &dist3, _DT_INDICES &idx1,
                                 _DT_INDICES &idx2, _DT_INDICES &idx3) {
    DeQueSimple(XYZ2);
    TBufGet(TEMP, _DT_XYZ2);
    TBufGet(WORK_LOCAL, _DT_XYZ2);
    print_tensor("XYZ1", xyz1, 8, "%f");
    print_tensor("XYZ2", LTensor(XYZ2), 40, "%f");
    auto newLoopCount = ALIGN_TO(vectorSize, 8);
    auto subLoopCount = CEIL_DIV(newLoopCount / 8, 255);
    print("vectorSize = %d\n", vectorSize);
    print("newLoopCount = %d\n", newLoopCount);
    print("subLoopCount = %d\n", subLoopCount);
    for (auto i = 0; i < subLoopCount; ++i) {
      auto repeatTimes = 255;
      if (i == subLoopCount - 1) {
        repeatTimes = (newLoopCount / 8) % 255;
      }
      print("repeatTimes = %d\n", repeatTimes);
      auto xyz2Index = i * elementsPerRepeat<_DT_XYZ1>() * 255;
      auto xyz2 = LTensor(XYZ2)[xyz2Index];
      Sub(xyz2, xyz1, xyz2, elementsPerRepeat<_DT_XYZ1>(), uint8_t(repeatTimes),
          {1, 0, 1, 8, 0, 8});
    }
    print_tensor("SUB XYZ2", LTensor(XYZ2), 40, "%f");
    Mul(BTensor(TEMP), LTensor(XYZ2), LTensor(XYZ2),
        vectorSize * elementsPerBlock<_DT_XYZ2>());
    print_tensor(u8"平方 TEMP", BTensor(TEMP), 40, "%f");
    auto repeatTimes = CEIL_DIV(vectorSize, 8);
    print("repeatTimes = %d\n", repeatTimes);
    constexpr uint64_t mask[2] = {0x0707070707070707, 0};
    BlockReduceSum(LTensor(XYZ2), BTensor(TEMP), repeatTimes, mask, 1, 1, 8);
    print_tensor(u8"和XYZ2", LTensor(XYZ2), 40, "%f");
    for (auto i = 0; i < 3; ++i) {
      ReduceMin(BTensor(TEMP), LTensor(XYZ2), BTensor(WORK_LOCAL), vectorSize,
                true);
      auto v = BTensor(TEMP).GetValue(0);
      auto vi = BTensor(TEMP).template ReinterpretCast<int32_t>().GetValue(1);
      print_tensor(u8"最小值TEMP", BTensor(TEMP), 8, "%f");
      tryBest(v, vi + mi, dist1, dist2, dist3, idx1, idx2, idx3);
      LTensor(XYZ2).SetValue(vi, _DT_XYZ2(1e40));
      print_tensor(u8"取消最小值XYZ2", LTensor(XYZ2), 40, "%f");
    }
    if (mi + vectorSize >= m) {
      auto index_1 = bi * n * 3 + ni * 3;
      GTensorName(DIST).SetValue(index_1, dist1);
      GTensorName(DIST).SetValue(index_1 + 1, dist2);
      GTensorName(DIST).SetValue(index_1 + 2, dist3);
      GTensorName(INDICES).SetValue(index_1, idx1);
      GTensorName(INDICES).SetValue(index_1 + 1, idx2);
      GTensorName(INDICES).SetValue(index_1 + 2, idx3);
    }
    QueFree(XYZ2);
  }

  __aicore__ inline void Process2() {
    auto mLoopCount = CEIL_DIV(m, tileVectorSize);
    auto finalVetorSize = m % tileVectorSize;
    for (auto bi = 0; bi < b; ++bi) {
      for (auto ni = 0; ni < n; ++ni) {
        {
          EnQueGlobal2Local(XYZ1, bi * n * 3 + ni * 3,
                            elementsPerBlock<_DT_XYZ1>());
        }
        {
          DeQueSimple(XYZ1);
          _DT_XYZ1 dist1 = 1e40, dist2 = 1e40, dist3 = 1e40;
          _DT_INDICES idx1 = -1, idx2 = -1, idx3 = -1;
          auto mi = 0;
          while (mi < m) {
            auto vectorSize = copyIn(bi, ni, mi);
            compute(LTensor(XYZ1), bi, ni, mi, vectorSize, dist1, dist2, dist3,
                    idx1, idx2, idx3);
            mi += vectorSize;
          }
          QueFree(XYZ1);
        }
      }
    }
  }
};

extern "C" __global__ __aicore__ void three_nn(GM_ADDR xyz1, GM_ADDR xyz2,
                                               GM_ADDR dist, GM_ADDR indices,
                                               GM_ADDR workspace,
                                               GM_ADDR tiling) {
  GET_TILING_DATA(tiling_data, tiling);
  Kernel<DTYPE_XYZ1, DTYPE_XYZ2, DTYPE_DIST, DTYPE_INDICES> op;
  op.Init(xyz1, xyz2, dist, indices, tiling_data);
  op.Process2();
}